package cn.xiaopengstack.domain.agent.model.valobj;

import cn.xiaopengstack.domain.agent.service.armory.element.RagAnswerAdvisor;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;

import java.util.HashMap;
import java.util.Map;

/**
 * @author jiangyangang
 */
@Getter
@AllArgsConstructor
@NoArgsConstructor
public enum AiClientAdvisorTypeEnumVO {
    CHAT_MEMORY("ChatMemory", "上下文记忆（内存模式）"){
        @Override
        public Advisor createAdvisor(AiClientAdvisorVO advisorVO, VectorStore vectorStore){
            AiClientAdvisorVO.ChatMemory chatMemory = advisorVO.getChatMemory();
            return PromptChatMemoryAdvisor.builder(
                    MessageWindowChatMemory.builder().maxMessages(chatMemory.getMaxMessages()).build()
            ).build();
        }
    },
    RAG_ANSWER("RagAnswer", "知识库"){
        @Override
        public Advisor createAdvisor(AiClientAdvisorVO advisorVO, VectorStore vectorStore){
            AiClientAdvisorVO.RagAnswer ragAnswer = advisorVO.getRagAnswer();
            return new RagAnswerAdvisor(vectorStore, SearchRequest.builder()
                    .topK(ragAnswer.getTopK())
                    .filterExpression(ragAnswer.getFilterExpression())
                    .build());
        }
    },
    ;

    private String type;
    private String info;

    // 静态Map缓存，用于快速查找
    private static final Map<String, AiClientAdvisorTypeEnumVO> TYPE_MAP = new HashMap<>();

    // 静态初始化块，在类加载时初始化Map
    static {
        for (AiClientAdvisorTypeEnumVO enumVO : values()) {
            TYPE_MAP.put(enumVO.getType(), enumVO);
        }
    }

    /**
     * 根据code获取枚举
     * @param type 编码
     * @return 枚举对象
     */
    public static AiClientAdvisorTypeEnumVO getByType(String type) {
        AiClientAdvisorTypeEnumVO enumVO = TYPE_MAP.get(type);
        if (enumVO == null) {
            throw new RuntimeException("err! advisorType " + type + " not exist!");
        }
        return enumVO;
    }


    public abstract Advisor createAdvisor(AiClientAdvisorVO advisorVO, VectorStore vectorStore);
}
